# Functions for the equal LRS model

# Initialises a matrix with n individuals
# individuals have a location (spX, spY) and are 
# initilised with gene for p.disp
init.inds<-function(n, spX, spY){
	X<-runif(n, 1, spX+1)%/%1 #grid ref
	Y<-runif(n, 1, spY+1)%/%1
	P<-sample(seq(0,1,0.1), n, replace=T) # dispersal genes
	cbind(X, Y, P)		
}

fec.table<-cbind(0:10, rep(3,11))

# reproduces based on table of fecundities and dispersal probs
# Global K
repro.glob<-function(popmatrix, spX, spY, fec.table, K, mutn, plot=T){
	dcat<-popmatrix[,"P"]*10
	dcat<-fec.table[pmatch(dcat, fec.table[,1], duplicates.ok=T), 2] #offspring per individual
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	density<-sum(density>0) #sum of occupied cells
	glob.size<-density*K
	juvs<-1:length(popmatrix[,1]) #adult indexes
	juvs<-rep(juvs, times=dcat)  #by number of offspring
	if (glob.size<length(juvs)) juvs<-sample(juvs, glob.size, replace=F) #sampled to K
	popmatrix<-popmatrix[juvs, ]
	mut.sample<-rbinom(length(juvs), 1, mutn) # individuals copping a mutation
	inds<-which(mut.sample==1)
	if (length(inds)>0){
		mutation<-sample(c(-0.2, -0.1, 0.1, 0.2), length(inds), replace=T, prob=c(0.1, 0.4, 0.4, 0.1))
		mutation<-popmatrix[inds, "P"]+mutation
		mutation[which(mutation>1)]<-1
		mutation[which(mutation<0.1)]<-0
		popmatrix[inds, "P"]<-mutation
	}
	if (plot==T) plotter(popmatrix, spX, spY, K)
	popmatrix		
}

#calculates neighbours on a grid excluding nonsense neighbours and returns a list of matrices of grid refs
#doesn't include source cell
neighbours.init<-function(spX, spY){
	neigh<-vector("list", length=spX*spY)
	for (x in 1:spX){
		for (y in 1:spY){
			out<-matrix(nrow=8, ncol=2)
			ifelse((x-1)==0, X<-NA, X<-x-1) #correct for indices going to zero
			ifelse((y-1)==0, Y<-NA, Y<-y-1)
			ifelse((y+1)>spY, Yp<-NA, Yp<-y+1) # correct for indices going to > space
			ifelse((x+1)>spX, Xp<-NA, Xp<-x+1)
			# Assign neighbours
			out[1,]<-c(Xp, y)
			out[2,]<-c(X, y)
			out[3,]<-c(x, Y)
			out[4,]<-c(Xp, Y)	
			out[5,]<-c(X, Y)
			out[6,]<-c(x, Yp)
			out[7,]<-c(Xp, Yp)
			out[8,]<-c(X, Yp)
			out<-subset(out, !is.na(apply(out, 1, sum)))
			neigh[[(y-1)*spX+x]]<-out
		}	
	}
	neigh	#return the list
}

# samples a matrix and returns one row of the matrix
matrix.sample<-function(mtrx, size=1){
	mtrx[sample(seq(length(mtrx[,1])), size=size),]
}

# executes nearest neighbour dispersal where individuals disperse to neighbouring cell with probability=prob.d 
# is called after reproduction/survival
disperse<-function(popmatrix, n.list, spX){
	if (length(popmatrix[,3])==0) return(popmatrix)
	disp<-which(rbinom(length(popmatrix[,1]), 1, popmatrix[,"P"])==1) #disperse or not
	sub.list<-n.list[(popmatrix[disp,"Y"]-1)*spX+popmatrix[disp,"X"]] #generate list of neighbour matrices
	sub.list<-lapply(sub.list, matrix.sample) #samples one row from each neighbour matrix
	sub.list<-do.call("rbind", sub.list) #places result into a matrix
	popmatrix[disp, 1:2]<-sub.list #assigns new locations back to popmatrix
	popmatrix
}

# plots population density and mean trait values through space
plotter<-function(popmatrix, spX, spY, K){
	density<-table(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)) # adult density through space
	density<-density/K
	trait<-tapply(popmatrix[,"P"], list(factor(popmatrix[,"X"], levels=1:spX), factor(popmatrix[,"Y"], levels=1:spY)), mean) #mean trait values through space
	par(mfrow=c(2,1))
	image(density, zlim=c(0,2), col=(grey((12:0)/12)), main="Population density")
	image(trait, zlim=c(0,1), col=(grey((12:0)/12)), main="Trait values")
}